﻿#if defined(__linux__)

#include	"Engine.Network.hpp"
#include	"Engine.Memory.hpp"
#include	"Engine.Thread.hpp"

#include	<atomic>
#include	<cstring>
#include	<errno.h>
#include	<map>

#include	<arpa/inet.h>
#include	<fcntl.h>
#include	<netinet/in.h>
#include	<sys/epoll.h>
#include	<sys/socket.h>
#include	<sys/types.h>
#include	<unistd.h>

namespace Engine { namespace Network {

	struct Session {
		int					nSocketFd;
		struct sockaddr_in	iAddr;
	};

	class SocketImpl : public Socket {
	public:
		SocketImpl();
		virtual ~SocketImpl();

		virtual int		Connect(const char * pStrIp, unsigned short uPort, IListener * pListener);
		virtual int		Send(const char * pData, size_t uSize);
		virtual void	Close(int nReason = ECloseReason::ClientClose);

	private:
		void			_IOThread();

	private:
		std::atomic<bool>	_bValid;
		size_t				_nRecvBufSize;
		IListener *			_pListener;
		Session				_iSession;
		Thread::Thread		_iIOWorker;
		char				_pReadBuff[ReceiveBuffer];
	};

	SocketImpl::SocketImpl()
		: _bValid(false)
		, _nRecvBufSize(0)
		, _pListener(nullptr)
		, _iSession()
		, _iIOWorker() {
	}

	SocketImpl::~SocketImpl() {
		Close();
	}

	int SocketImpl::Connect(const char * pStrIp, unsigned short uPort, IListener * pListener) {
		if (!pStrIp || strlen(pStrIp) <= 1 || uPort <= 0 || !pListener) return ErrParam;
		if (_bValid) return ErrAlreadyUsed;

		if ((_iSession.nSocketFd = ::socket(AF_INET, SOCK_STREAM, 0)) < 0)
			return errno;

		int nFlags = fcntl(_iSession.nSocketFd, F_GETFL);
		fcntl(_iSession.nSocketFd, F_SETFL, O_NONBLOCK | nFlags);

		::memset(&_iSession.iAddr, 0, sizeof(_iSession.iAddr));
		_iSession.iAddr.sin_family		= AF_INET;
		_iSession.iAddr.sin_port		= htons(uPort);
		_iSession.iAddr.sin_addr.s_addr	= inet_addr(pStrIp);

		if (::connect(_iSession.nSocketFd, (sockaddr *) &_iSession.iAddr, sizeof(_iSession.iAddr)) < 0) {
			if (errno != EINPROGRESS) {
				return errno;
			} else {
				int nChecker = ::epoll_create(1);
				if (nChecker == -1) {
					::close(_iSession.nSocketFd);
					return errno;
				}

				struct epoll_event iEvent;
				iEvent.data.fd	= _iSession.nSocketFd;
				iEvent.events	= EPOLLIN | EPOLLOUT | EPOLLET;

				if (::epoll_ctl(nChecker, EPOLL_CTL_ADD, _iSession.nSocketFd, &iEvent) != 0) {
					::close(nChecker);
					::close(_iSession.nSocketFd);
					return errno;
				}

				while (true) {
					::epoll_wait(nChecker, &iEvent, 1, 3000);
					if (iEvent.data.fd == _iSession.nSocketFd) {
						int nStatus = 0;
						size_t nLen = sizeof(nStatus);
						if (::getsockopt(_iSession.nSocketFd, SOL_SOCKET, SO_ERROR, (void *) &nStatus, (socklen_t *) &nLen) < 0) {
							::epoll_ctl(nChecker, EPOLL_CTL_DEL, _iSession.nSocketFd, NULL);
							::close(nChecker);
							return -1;
						} else if (nStatus != 0) {
							::epoll_ctl(nChecker, EPOLL_CTL_DEL, _iSession.nSocketFd, NULL);
							::close(nChecker);
							return ErrTimeOut;
						} else {
							break;
						}
					}
                }				
			}
		}

		size_t nBuffLen = sizeof(_nRecvBufSize);
		if (::getsockopt(_iSession.nSocketFd, SOL_SOCKET, SO_RCVBUF, &_nRecvBufSize, (socklen_t *)&nBuffLen) < 0) {
			::close(_iSession.nSocketFd);
			return errno;
		}
		
		_bValid				= true;
		_pListener			= pListener;
		_pListener->pSender	= this;

		if (!_iIOWorker.Start(&SocketImpl::_IOThread, this)) {
			_bValid = false;
			::close(_iSession.nSocketFd);
			return ErrCreateThread;
		}

		_pListener->OnConnect(_iSession.iAddr.sin_addr.s_addr, uPort);
        return Success;
	}

	int SocketImpl::Send(const char * pData, size_t uSize) {
		char *	pSend	= (char *)pData;
		ssize_t	nSend	= 0;
		size_t	nTotal	= uSize;

        while (true) {
			nSend = ::send(_iSession.nSocketFd, pSend, nTotal, MSG_DONTWAIT);
			if (nSend < 0) {
				if (errno == EAGAIN) {
					::usleep(1000);
				} else {
					return errno;
				}
			} else if ((size_t) nSend < nTotal){
				nTotal	-= nSend;
				pSend	+= nSend;
			} else {
				break;
			}
        }

		return Success;
	}

	void SocketImpl::Close(int nReason) {
		if (!_bValid) return;
		_bValid = false;

		::close(_iSession.nSocketFd);
		_pListener->OnClose(nReason);
	}

	void SocketImpl::_IOThread() {
		while (_bValid) {
			memset(_pReadBuff, 0, ReceiveBuffer);

			ssize_t nReaded		= 0;
			ssize_t nCurRead	= 0;

			while (nReaded < ReceiveBuffer) {
				nCurRead = ::recv(_iSession.nSocketFd, _pReadBuff + nReaded, ReceiveBuffer - nReaded, MSG_DONTWAIT);
				if (nCurRead == 0) {
					Close(ECloseReason::SystemClose);
					break;
				} else if (nCurRead < 0) {
					if (errno != EAGAIN) Close(ECloseReason::DataError);
					break;
				} else if (nCurRead < _nRecvBufSize) {
					nReaded += nCurRead;
					break;
				} else {
					nReaded += nCurRead;
				}
			}

			if (nReaded > 0) _pListener->OnReceive(_pReadBuff, nReaded);
		}
	}

	class ServerSocketImpl : public ServerSocket {
	public:
		ServerSocketImpl();
		virtual ~ServerSocketImpl();

		virtual int		Listen(const char * pStrIp, unsigned short uPort, IListener * pListener);
		virtual int		Send(uint64_t nConnId, const char * pData, size_t uSize);
		virtual int		Broadcast(const char * pData, size_t uSize);
		virtual void	Close(uint64_t nConnId, int nReason = ECloseReason::SystemClose);
		virtual void	Shutdown();

	private:
		void		_AcceptThread();
		void		_IOThread();

	private:
		std::atomic<bool>			_bValid;
		uint64_t					_nNextConnId;
		size_t						_nRecvBufSize;
		IListener *					_pListener;
		int							_nEpollAcceptFd;
		int							_nEpollIOFd;
		Session						_iSession;
		std::map<uint64_t, Session>	_mpAllClient;
		Thread::Thread				_iAcceptWorker;
		Thread::Thread				_iIOWorker;
		char						_pReadBuff[ReceiveBuffer];
	};

	ServerSocketImpl::ServerSocketImpl()
		: _bValid(false)
		, _nNextConnId(0)
		, _nRecvBufSize(0)
		, _pListener(nullptr)
		, _nEpollAcceptFd(-1)
		, _nEpollIOFd(-1)
		, _iSession()
		, _mpAllClient()
		, _iAcceptWorker()
		, _iIOWorker() {
	}

	ServerSocketImpl::~ServerSocketImpl() {
		Shutdown();
	}

	int ServerSocketImpl::Listen(const char * pStrIp, unsigned short uPort, IListener * pListener) {
		if (!pStrIp || strlen(pStrIp) <= 0 || uPort <= 0 || !pListener) return ErrParam;
		if (_bValid) return ErrAlreadyUsed;

		if ((_nEpollAcceptFd = ::epoll_create(1)) == -1) return errno;
		if ((_nEpollIOFd = ::epoll_create(1)) == -1) return errno;

		if ((_iSession.nSocketFd = ::socket(AF_INET, SOCK_STREAM, 0)) < 0) {
			::close(_nEpollAcceptFd);
			::close(_nEpollIOFd);
			return errno;
		}

		int nFlags = fcntl(_iSession.nSocketFd, F_GETFL);
		fcntl(_iSession.nSocketFd, F_SETFL, O_NONBLOCK | nFlags);

		size_t nBufLen = sizeof(_nRecvBufSize);
		if (::getsockopt(_iSession.nSocketFd, SOL_SOCKET, SO_RCVBUF, &_nRecvBufSize, (socklen_t *)&nBufLen) < 0) {
			::close(_nEpollAcceptFd);
			::close(_nEpollIOFd);
			::close(_iSession.nSocketFd);
			return errno;
		}

		struct epoll_event iEvent;
		iEvent.data.fd	= _iSession.nSocketFd;
		iEvent.events	= EPOLLIN | EPOLLOUT | EPOLLET;

		if (epoll_ctl(_nEpollAcceptFd, EPOLL_CTL_ADD, _iSession.nSocketFd, &iEvent) != 0) {
			::close(_nEpollAcceptFd);
			::close(_nEpollIOFd);
			::close(_iSession.nSocketFd);
			return errno;
		}

        ::memset(&_iSession.iAddr, 0, sizeof(_iSession.iAddr));
        _iSession.iAddr.sin_family		= AF_INET;
        _iSession.iAddr.sin_port		= htons(uPort);
        _iSession.iAddr.sin_addr.s_addr	= inet_addr(pStrIp);

        if (::bind(_iSession.nSocketFd, (sockaddr *)&_iSession.iAddr, sizeof(_iSession.iAddr)) < 0) {
			::close(_nEpollAcceptFd);
			::close(_nEpollIOFd);
			::close(_iSession.nSocketFd);
			return errno;
        }

        if (::listen(_iSession.nSocketFd, 512) < 0) {
            ::close(_nEpollAcceptFd);
			::close(_nEpollIOFd);
            ::close(_iSession.nSocketFd);
            return errno;
        }

		_bValid				= true;
		_pListener			= pListener;
		_pListener->pSender = this;

		if (!_iAcceptWorker.Start(&ServerSocketImpl::_AcceptThread, this) || !_iIOWorker.Start(&ServerSocketImpl::_IOThread, this)) {
			_bValid = false;
			::close(_nEpollAcceptFd);
			::close(_nEpollIOFd);
            ::close(_iSession.nSocketFd);
            return ErrCreateThread;
		}

		return Success;
	}

	int ServerSocketImpl::Send(uint64_t nConnId, const char * pData, size_t uSize) {
		if (!_bValid || nConnId > _nNextConnId || nConnId <= 0 || _mpAllClient.find(nConnId) == _mpAllClient.end() || !pData || uSize <= 0)
			return ErrParam;

		Session	iClient	= _mpAllClient[nConnId];
		char *	pSend	= (char *)pData;
		ssize_t	nSend	= 0;
		size_t	nTotal	= uSize;

        while (true) {
			nSend = ::send(iClient.nSocketFd, pSend, nTotal, MSG_DONTWAIT);
			if (nSend < 0) {
				if (errno == EAGAIN) {
					::usleep(1000);
				} else {
					return errno;
				}
			} else if ((size_t) nSend < nTotal){
				nTotal	-= nSend;
				pSend	+= nSend;
			} else {
				break;
			}
        }

		return Success;
	}

	int ServerSocketImpl::Broadcast(const char * pData, size_t uSize) {
		if (!_bValid)
			return ErrInvalidConnId;

		for (auto itor = _mpAllClient.begin(); itor != _mpAllClient.end(); ++itor) {
			Session	iClient	= itor->second;
			char *	pSend	= (char *)pData;
			ssize_t	nSend	= 0;
			size_t	nTotal	= uSize;

			while (true) {
				nSend = ::send(iClient.nSocketFd, pSend, nTotal, MSG_DONTWAIT);
				if (nSend < 0) {
					if (errno == EAGAIN) {
						::usleep(1000);
					} else {
						break;	//! 注意，这里没有处理错误
					}
				} else if ((size_t) nSend < nTotal){
					nTotal	-= nSend;
					pSend	+= nSend;
				} else {
					break;
				}
			}
		}

		return Success;
	}

	void ServerSocketImpl::Close(uint64_t nConnId, int nReason) {
		if (nConnId > _nNextConnId)
			return;

		auto itor = _mpAllClient.find(nConnId);
		if (itor == _mpAllClient.end())
			return;

		_pListener->OnClose(nConnId, nReason);

		::epoll_ctl(_nEpollIOFd, EPOLL_CTL_DEL, itor->second.nSocketFd, NULL);
		::close(itor->second.nSocketFd);
		_mpAllClient.erase(itor);
	}

	void ServerSocketImpl::Shutdown() {
		if (!_bValid) return;
		_bValid = false;

		for (auto itor = _mpAllClient.begin(); itor != _mpAllClient.end(); ++itor)
			Close(itor->first, ECloseReason::SystemClose);

		_mpAllClient.clear();

		::close(_nEpollAcceptFd);
		::close(_nEpollIOFd);
        ::close(_iSession.nSocketFd);
	}

	void ServerSocketImpl::_AcceptThread() {
		struct epoll_event	lEvents[512];
		struct epoll_event	iEvent;
		struct sockaddr_in	iClientAddr;

		socklen_t			nSize = sizeof(iClientAddr);

		iEvent.events = EPOLLIN | EPOLLET | EPOLLOUT;

		while (_bValid) {
			int nEvents = ::epoll_wait(_nEpollAcceptFd, lEvents, 512, 50);

			for (int nIdx = 0; nIdx < nEvents; ++nIdx) {
				if (lEvents[nIdx].data.fd == _iSession.nSocketFd) {
					int iClient = ::accept(_iSession.nSocketFd, (struct sockaddr*)(&iClientAddr), &nSize);
					if (iClient < 0) continue;

					uint64_t nConnId = ++_nNextConnId;
					iEvent.data.u64 = nConnId;

					if (epoll_ctl(_nEpollIOFd, EPOLL_CTL_ADD, iClient, &iEvent) != 0) {
						::close(iClient);
						_pListener->OnAccept(nConnId, iClientAddr.sin_addr.s_addr, iClientAddr.sin_port, errno);
						continue;
					}

					Session iClientSession;
					iClientSession.nSocketFd	= iClient;
					iClientSession.iAddr		= iClientAddr;
					_mpAllClient[nConnId]		= iClientSession;

					_pListener->OnAccept(nConnId, iClientAddr.sin_addr.s_addr, iClientAddr.sin_port, Success);
				}
			}
		}
	}

	void ServerSocketImpl::_IOThread() {
		struct epoll_event	lEvents[512];

		while (_bValid) {
			int nEvents = ::epoll_wait(_nEpollIOFd, lEvents, 512, 50);

			for (int nIdx = 0; nIdx < nEvents; ++nIdx) {
				if (lEvents[nIdx].events & EPOLLIN) {
					uint64_t nConnId = lEvents[nIdx].data.u64;

					auto itor = _mpAllClient.find(nConnId);
					if (itor == _mpAllClient.end()) continue;

					::memset(_pReadBuff, 0, ReceiveBuffer);

					ssize_t nReaded		= 0;
					ssize_t nCurRead	= 0;

					while (nReaded < ReceiveBuffer) {
						nCurRead = ::recv(itor->second.nSocketFd, _pReadBuff + nReaded, ReceiveBuffer - nReaded, MSG_DONTWAIT);
						if (nCurRead == 0) {
							Close(nConnId, ECloseReason::SystemClose);
							break;
						} else if (nCurRead < 0) {
							if (errno != EAGAIN) Close(nConnId, ECloseReason::DataError);
							break;
						} else if (nCurRead < _nRecvBufSize) {
							nReaded += nCurRead;
							break;
						} else {
							nReaded += nCurRead;
						}
					}

					if (nReaded > 0) _pListener->OnReceive(nConnId, _pReadBuff, nReaded);
				}
			}
		}
	}

	Socket * Socket::Create() {
		return new (std::nothrow) SocketImpl();
	}

	void Socket::Close(Socket * pClient, int nReason) {
		if (!pClient) return;
		SocketImpl * pConn = (SocketImpl *)pClient;
		delete pConn;
	}

	ServerSocket * ServerSocket::Create() {
		return new (std::nothrow) ServerSocketImpl();
	}

	void ServerSocket::Shutdown(ServerSocket * pServer) {
		if (!pServer) return;
		ServerSocketImpl * pImpl = (ServerSocketImpl *) pServer;
		delete pImpl;
	}


}}

#endif

